# This file is part of MaixPY
# Copyright (c) sipeed.com
#
# Licensed under the MIT license:
#   http://www.opensource.org/licenses/mit-license.php

import socket
import threading
import time

import numpy as np
import torch
import os

local_ip = "172.20.10.3"
local_port = 8888


def xyxy2xywh(x):
    """
    将锚框从 nx4 boxes([x1, y1, x2, y2]) 格式转换成 maixpy 绘图所需的格式 [x1, y1, w, h]
    xy1=top-left, xy2=bottom-right
    """
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = x[:, 0]  # x1=top-left.x
    y[:, 1] = x[:, 1]  # y1=top-left.y
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y.int()  # maixpy image.draw_rectangle 暂只支持 int


def getPre(model, img, imgSize):
    """
    参数：
    model：yolo5模型
    img：图像路径
    imgSize：方形图像大小，如 200
    返回结果示例(patches_253.jpg)：
    patches 0 28 29 64
    patches 0 145 12 54
    patches 63 61 66 91
    patches 129 61 36 54
    """
    results = model(img, size=imgSize)  # 指定图像大小，不指定的效果很差
    pre = results.pandas().xyxy[0].sort_values("xmin")
    tmp = pre.iloc[:, 0:4]
    temp = np.array(tmp).reshape(-1, 4)
    xywh = (xyxy2xywh(torch.tensor(temp))).tolist()  # normalized xywh
    lines = ""
    for i in range(len(pre)):
        tmpres = pre.loc[i, "name"], " ".join(map(str, xywh[i]))
        lines += (
            (str(tmpres) + ("\n" if i < len(pre) - 1 else ""))
            .replace(",", "")
            .replace("'", "")
            .replace("(", "")
            .replace(")", "")
        )
    return lines


def server():
    while True:
        conn, addr = sk.accept()
        print("hello client,ip:", addr)
        t = threading.Thread(target=receiveThread, args=(conn,))
        t.daemon = True
        t.start()


def receiveThread(conn):
    # conn.settimeout(10)
    conn_end = False
    pack_size = 100000  # 10w
    print("begin to receive")
    while True:
        img = b""
        client_data = None
        errTime = 0  # 错误计数
        if conn_end:
            break
        while True:
            try:
                client_data = conn.recv(2048)
                # print("rcv:", len(client_data))
                time.sleep(0.015)
            except Exception(e):
                print(e)
                break
            if not client_data:
                # print("received null")
                errTime += 1
                if errTime > 299:  # 多次收到空包，认定为结束
                    conn_end = True
                    break
                continue
            else:
                img += client_data
            if client_data.endswith(b"\xFF\xD9"):  # jpg 尾部检验
                print("image end dect. recive end, pic len:", len(img))
                f = open("tmp.jpg", "wb")  # 写入图片
                f.write(img)
                f.close()
                errTime = 0  # 错误次数清零

                res = getPre(MODEL, "tmp.jpg", 200)  # 使用模型进行目标检测
                print(res)

                conn.send(res.encode("utf-8"))
                break
            if len(img) > pack_size:  # 数据异常大
                print("too big, unabnomal, clear")
                conn_end = True
                break
    conn.close()
    print("receive thread end")


# 加载模型
MODEL = torch.hub.load(
    "../yolov5",
    "custom",
    path="../yolov5/runs/train/exp11/weights/best",
    source="local",
)
# img = r"/run/media/kearney/a/CAU/42course/毕设/MaxiPy/tmp.jpg"
# results = MODEL(img, size=200)
# results.print()
# results.save(os.getcwd())  # 在当前目录下保存图像，同名文件会被覆盖
ip_port = (local_ip, local_port)
sk = socket.socket()
sk.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sk.bind(ip_port)
sk.listen(50)
print("accept now,wait for client")


tmp = threading.Thread(target=server, args=())
tmp.daemon = True
tmp.start()

while True:
    pass
